from time import sleep

import torch
import matplotlib.pyplot as plt


# 定义一个函数，作为目标函数
def f(x):
    return (2 / 5) * x ** 5 + (3 / 4) * x ** 4 - 2 * x ** 3 - x ** 2 + 3 * x + 1


if __name__ == '__main__':
    # 生成数据
    x = torch.linspace(start=-10, end=10, steps=100)
    y = f(x)

    # 使用pytorch初始化一个随机函数, 并且需要求导
    a = torch.tensor(1.0, requires_grad=True)
    b = torch.tensor(1.0, requires_grad=True)
    c = torch.tensor(1.0, requires_grad=True)
    d = torch.tensor(1.0, requires_grad=True)
    e = torch.tensor(1.0, requires_grad=True)
    f = torch.tensor(1.0, requires_grad=True)

    # 定义优化器
    optimizer = torch.optim.Adam(params=[a, b, c, d, e, f], lr=0.001)

    # 定义损失函数
    loss_fn = torch.nn.MSELoss()

    # 开始训练
    for epoch in range(1000000):
        y_pred = a * x ** 5 + b * x ** 4 + c * x ** 3 + d * x ** 2 + e * x + f
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 每隔 1 秒展示一次函数与目标函数的差距
        if epoch % 10000 == 0:
            with torch.no_grad():
                # 绘制数据
                plt.plot(x, y, label='y_true')
                # 绘制预测函数
                plt.plot(x, y_pred.detach().numpy(), label='y_pred')
                # 指定x轴和y轴的范围
                plt.xlim(-10, 10)
                plt.ylim(-10, 10)
                plt.legend()
                # 保存图片
                plt.savefig(str(epoch) + '.png')
                plt.show()

